import os
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader

# 多分辨率生成（固定4种分辨率）
RESOLUTIONS = [(256, 256), (128, 128), (64, 64), (32, 32)]

# 无监督数据增强（用于生成正例）
class SimCLRDataTransform:
    def __init__(self, resolution):
        self.transform = transforms.Compose([
            # 随机裁剪并调整大小至指定分辨率，增加样本多样性
            transforms.RandomResizedCrop(resolution),
            # 以0.5的概率随机水平翻转，模拟不同视角
            transforms.RandomHorizontalFlip(p=0.5),
            # 以0.8的概率应用颜色抖动，参数分别为亮度、对比度、饱和度和色调的调整范围
            transforms.RandomApply([transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),
            # 以0.2的概率转换为灰度图，保留图像结构信息
            transforms.RandomGrayscale(p=0.2),
            # 应用高斯模糊，核大小根据分辨率动态调整，sigma在0.1到2.0之间随机
            transforms.GaussianBlur(kernel_size=int(0.1 * resolution[0]) + 1, sigma=(0.1, 2.0)),
            # 将PIL图像转换为PyTorch张量，范围[0,1]
            transforms.ToTensor(),
            # 使用ImageNet预训练模型的均值和标准差进行归一化
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __call__(self, x):
        return [self.transform(x), self.transform(x)]  # 生成2个增强视图

# 数据集类（融合多分辨率和增强样本）
class MultiResDataset(Dataset):
    def __init__(self, root_dir, split='train', positive_pairs=None):
        self.root_dir = root_dir
        self.split = split
        self.classes = os.listdir(root_dir)
        self.class_to_idx = {cls: i for i, cls in enumerate(self.classes)}
        self.positive_pairs = positive_pairs  # 正例对列表，格式为[(img_path1, img_path2), ...]
        self.samples = self._load_samples()
        self.base_transform = transforms.ToTensor()

    def _load_samples(self):
        # todo: @ningzhiyuan 结合实际数据集进行修改，这里需要重点处理正例怎么处理，以及获得正确的labels
        # 支持输入正例并生成labels
        samples = []
        
        # 如果提供了正例对，则使用这些正例对
        if self.positive_pairs is not None:
            # 为每对正例分配相同的label
            for pair_id, (img_path1, img_path2) in enumerate(self.positive_pairs):
                # 确保路径存在
                if os.path.exists(img_path1) and os.path.exists(img_path2):
                    samples.append((img_path1, pair_id))
                    samples.append((img_path2, pair_id))
        else:
            # 否则使用原始的按文件夹分类方式
            for cls in self.classes:
                cls_dir = os.path.join(self.root_dir, cls)
                if not os.path.isdir(cls_dir):
                    continue  # 跳过非目录文件
                
                for img_name in os.listdir(cls_dir):
                    img_path = os.path.join(cls_dir, img_name)
                    # 存储原始图片路径+类别ID
                    samples.append((img_path, self.class_to_idx[cls]))
        
        return samples

    def __getitem__(self, idx):
        # todo: @ningzhiyuan 结合实际数据集进行修改，
        img_path, label = self.samples[idx]
        img = Image.open(img_path).convert('RGB')
        
        # 生成多分辨率版本
        multi_res_imgs = []
        for res in RESOLUTIONS:
            resized = img.resize(res, Image.BILINEAR)
            multi_res_imgs.append(resized)
        
        # 对每个分辨率生成增强样本（训练时）
        if self.split == 'train':
            augmented = []
            for res_img in multi_res_imgs:
                transform = SimCLRDataTransform(res_img.size)
                aug1, aug2 = transform(res_img)
                augmented.extend([aug1, aug2])
            return {'images': augmented, 'label': label}
        else:
            # 测试时只返回原始分辨率特征
            img_tensor = self.base_transform(img.resize(RESOLUTIONS[0]))
            return {'image': img_tensor, 'label': label}

    def __len__(self):
        return len(self.samples)

# 数据加载器
def get_dataloader(root_dir, batch_size=32, split='train', positive_pairs=None):
    dataset = MultiResDataset(root_dir, split, positive_pairs)
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=(split == 'train'),
        num_workers=4,
        pin_memory=True
    )